//
// Created by xuwj on 2021/11/1.
//

#ifndef TRACK_ASSISTANT_VISUAL_FRAMES_SYSTEM_H
#define TRACK_ASSISTANT_VISUAL_FRAMES_SYSTEM_H

#include "tf_publisher.h"
#include "tf/transform_listener.h"
#include "geometry_msgs/PoseStamped.h"
#include "geometry_msgs/PoseWithCovarianceStamped.h"

#include "memory"
#include "mutex"
#include "Eigen/Core"
#include "Converter.h"
#include "TagItems.h"

using namespace std;
namespace VNavAssistant{

    /** A system to manage tf frames that related to visual-nav
     *  HERE, The visual-pose-publisher will manage tag_odom and it's visual-pose by itself
     *  the structure of the system is:
     *  (parent->child)  map -> odom_robot -> base_footprint -> base_link -> camera_link
     */
    class PoseRectifySystem{
    public:
        PoseRectifySystem(const string cam_f, const string odom_f, const string map_f = "/map",
                 const string base_f = "/base_link", const string base_footprint_f = "/base_footprint");

        //用于 与会利用tag信息做位姿修正的ORB 配合 的CallBack函数
        void SLAMPoseCallback(const geometry_msgs::PoseStamped::ConstPtr &msg);

    private:
        //frame_id
        string map_frame, odom_frame, base_footprint_frame, base_link_frame, camera_frame;

        TFPublisher* pTF_pub;
        tf::TransformListener listener;
    };


    struct PoseWithTime{
        PoseWithTime(){};
        PoseWithTime(const cv::Mat T, const double t):
        T_cam0_camk_(T), timestamp_(t){}

        cv::Mat T_cam0_camk_;
        double timestamp_;
    };

    struct PoseManager{

        PoseManager(const cv::Mat T_camk_map_, const cv::Mat T_cam0_camk_, const double time);

        cv::Mat T_camk_map, T_cam0_camk;
        double timestamp;

        cv::Mat T_map_cam0;
        Eigen::Quaterniond quat;
        Eigen::Vector3d t;

        bool is_good_;
    };

    //用于统计 接近的位姿 的平均值， 以便对位姿分组统计直方图
    struct PoseStats{

        PoseStats(){
            for(int i=0; i<4; ++i)
                avg_q[i] = 0.0;
            for(int i=0; i<3; ++i)
                avg_t[i] = 0.0;
            counts = 0;
        }

        void push_back(const float* q, const float* t);
        void push_back(const Eigen::Quaterniond &q, const Eigen::Vector3d &t);

        float avg_q[4];   // x, y, z, w
        float avg_t[3];
        int counts = 0;
    };

    class PoseStatisticalGrouping{
    public:
        PoseStatisticalGrouping(const float th_Q_xyz, const float th_Q_w, const float th_t_xy, const float th_t_z);

        void dealWith(int index, const PoseManager &pm);

    public:
        vector<PoseStats> vGroupPoses;
        vector<vector<int>> vGroupIndices;
        int max_index, max_counts;
        int second_max_index, second_max_counts;

        float th_rotation_xyz, th_rotation_w;
        float th_translation_xy, th_translation_z;
    };

    /** A system to manage tf frames that related to visual-nav
     *  HERE, The visual-pose-publisher only publish the relative-pose: T_cam0_camk
     *  the structure of the system is:
     *  (parent->child)  map -> odom_robot -> base_link -> camera_link
     */
    class VisualOdomTFManager{
    public:
        VisualOdomTFManager(ros::NodeHandle &nh, const string config_path, const string tag_topic, int buff_size, const string cam_f, const string odom_f,
                            const string map_f = "/map", const string base_f = "/base_link", const string base_footprint_f = "/base_footprint");

        // 处理来自ORB-SLAM2的每一帧位姿Tcam0_camk
        void OrbOdomPoseCallback(const geometry_msgs::PoseStamped::ConstPtr &msg);

        // 处理来自Vins-Mono的每一帧位姿
        void VinsOdomPoseCallback(const nav_msgs::Odometry::ConstPtr &msg);

        // 处理来自initial_pose的位姿
        void InitialPoseCallback(const geometry_msgs::PoseWithCovarianceStamped::ConstPtr &msg);

        void MainProcess(const cv::Mat T_cam0_camk, const double timestamp, ros::Time tstamp);

        bool CheckDataQueue();
        void Initialize();
        void CorrectPose();

    public:
        ros::Subscriber odom_subscriber_;
        ros::Subscriber intial_pose_subscriber_;

    private:
        //frame_id
        string map_frame, odom_frame, base_footprint_frame, base_link_frame, camera_frame;

        //TF
//        TFPublisher* pTF_pub;
        tf::TransformBroadcaster* pTF_pub;
        tf::TransformListener listener;

        // AprilTag
        string tag_topic_;
        deque<TagItems::TagOdomData> tagOdom_dataBuf_;
        shared_ptr<TagItems::OdomSubscriber> pTagPoseSubscriber;

        // Visual-Localization algorithm related
        deque<PoseWithTime> visualPose_dataBuf_;
        vector<PoseManager> mvNeedProcessData_;

        bool initialized;
        cv::Mat mT_map_cam0;
        cv::Mat curT_cam0_camk;
        std::mutex mutexPoseCam;
        std::mutex mutexPoseMap;
        tf::Transform trans_map_cam0;
        double last_corrected_time;
        double cur_time;
        int stuck_count;

        // Configs
        double trans_time_tolerance;

    };


}

#endif //TRACK_ASSISTANT_VISUAL_FRAMES_SYSTEM_H
